import os
import time
import tensorflow as tf
import pandas as pd
import matplotlib.pyplot as plt
from tensorflow.keras.models import save_model


from net.cow_net import ConvModel
from utils.data_generator import train_val_generator
from utils.image_plot import plot_images

#读取训练集
train_gen = train_val_generator(
    data_dir='../data/imagenette2-160/train',
    target_size=(64, 64),
    batch_size=32,
    class_mode='categorical',
    subset='training')

val_gen = train_val_generator(
    data_dir='../data/imagenette2-160/train',
    target_size=(64, 64),
    batch_size=32,
    class_mode='categorical',
    subset='validation')

class_names = list(train_gen.class_indices.keys())
print(class_names)
#ImageDataGenerator的返回结果是个迭代器，
#取15张训练集查看
train_batch, train_label_batch = train_gen.next()
#plot_images(train_batch, train_label_batch)
plot_images(train_batch, train_label_batch, class_names)

val_batch, val_label_batch = val_gen.next()
#plot_images(val_batch, val_label_batch)
plot_images(val_batch, val_label_batch, class_names)

#实例化

model = ConvModel()


#设置损失函数loss, optimizer metrics
model.compile(
    loss='categorical_crossentropy',
    optimizer=tf.keras.optimizers.SGD(
        learning_rate=0.001),
    metrics=['accuracy']
)


#train model
"""history = model.fit(
    x = train_gen, steps_per_epoch=351,epochs= 100,
    validation_data=val_gen,validation_steps=88,shuffle=True)
"""
history = model.fit(
    x = train_gen, steps_per_epoch=351, epochs=100,
    validation_data=val_gen, validation_steps=88, shuffle=True)



#show history
pd.DataFrame(history.history).plot(figsize=(8, 5))
plt.grid(True)
plt.xlabel('epoch')
plt.show()


#save model
model_name = "model-" + time.strftime('%Y-%m-%d-%H-%M-%S')
model_path = os.path.join('..', 'models', model_name)
if not os.path.exists(model_path):
    os.makedirs(model_path)

    # save model
save_model(model = model, filepath= model_path)



